// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Array accessor classes for Binary, LargeBinary, String, LargeString,
// FixedSizeBinary

#pragma once

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <vector>

#include <nebula/array/array_base.h>
#include <nebula/array/data.h>
#include <nebula/core/buffer.h>
#include <nebula/core/stl_iterator.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <turbo/base/macros.h>

namespace nebula {

    /// \addtogroup binary-arrays
    ///
    /// @{

    // ----------------------------------------------------------------------
    // Binary and String

    /// Base class for variable-sized binary arrays, regardless of offset size
    /// and logical interpretation.
    template<typename TYPE>
    class BaseBinaryArray : public FlatArray {
    public:
        using TypeClass = TYPE;
        using offset_type = typename TypeClass::offset_type;
        using IteratorType = stl::ArrayIterator<BaseBinaryArray<TYPE>>;

        /// Return the pointer to the given elements bytes
        // XXX should get_value(int64_t i) return a string_view?
        const uint8_t *get_value(int64_t i, offset_type *out_length) const {
            // Account for base offset
            i += data_->offset;
            const offset_type pos = raw_value_offsets_[i];
            *out_length = raw_value_offsets_[i + 1] - pos;
            return raw_data_ + pos;
        }

        /// \brief Get binary value as a string_view
        ///
        /// \param i the value index
        /// \return the view over the selected value
        [[nodiscard]] std::string_view get_view(int64_t i) const {
            // Account for base offset
            i += data_->offset;
            const offset_type pos = raw_value_offsets_[i];
            return std::string_view(reinterpret_cast<const char *>(raw_data_ + pos),
                                    raw_value_offsets_[i + 1] - pos);
        }

        std::optional<std::string_view> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        /// \brief Get binary value as a string_view
        /// Provided for consistency with other arrays.
        ///
        /// \param i the value index
        /// \return the view over the selected value
        [[nodiscard]] std::string_view value(int64_t i) const { return get_view(i); }

        /// \brief Get binary value as a std::string
        ///
        /// \param i the value index
        /// \return the value copied into a std::string
        [[nodiscard]] std::string get_string(int64_t i) const { return std::string(get_view(i)); }

        /// Note that this buffer does not account for any slice offset
        [[nodiscard]] std::shared_ptr<Buffer> value_offsets() const { return data_->buffers[1]; }

        /// Note that this buffer does not account for any slice offset
        [[nodiscard]] std::shared_ptr<Buffer> value_data() const { return data_->buffers[2]; }

        [[nodiscard]] const offset_type *raw_value_offsets() const {
            return raw_value_offsets_ + data_->offset;
        }

        [[nodiscard]] const uint8_t *raw_data() const { return raw_data_; }

        /// \brief Return the data buffer absolute offset of the data for the value
        /// at the passed index.
        ///
        /// Does not perform boundschecking
        [[nodiscard]] offset_type value_offset(int64_t i) const {
            return raw_value_offsets_[i + data_->offset];
        }

        /// \brief Return the length of the data for the value at the passed index.
        ///
        /// Does not perform boundschecking
        [[nodiscard]] offset_type value_length(int64_t i) const {
            i += data_->offset;
            return raw_value_offsets_[i + 1] - raw_value_offsets_[i];
        }

        /// \brief Return the total length of the memory in the data buffer
        /// referenced by this array. If the array has been sliced then this may be
        /// less than the size of the data buffer (data_->buffers[2]).
        [[nodiscard]] offset_type total_values_length() const {
            if (data_->length > 0) {
                return raw_value_offsets_[data_->length + data_->offset] -
                       raw_value_offsets_[data_->offset];
            } else {
                return 0;
            }
        }

        [[nodiscard]] IteratorType begin() const { return IteratorType(*this); }

        [[nodiscard]]  IteratorType end() const { return IteratorType(*this, length()); }

    protected:
        // For subclasses
        BaseBinaryArray() = default;

        // Protected method for constructors
        void set_data(const std::shared_ptr<ArrayData> &data) {
            this->Array::set_data(data);
            raw_value_offsets_ = data->GetValuesSafe<offset_type>(1, /*offset=*/0);
            raw_data_ = data->GetValuesSafe<uint8_t>(2, /*offset=*/0);
        }

        const offset_type *raw_value_offsets_ = nullptr;
        const uint8_t *raw_data_ = nullptr;
    };

    /// Concrete Array class for variable-size binary data
    class TURBO_EXPORT BinaryArray : public BaseBinaryArray<BinaryType> {
    public:
        explicit BinaryArray(const std::shared_ptr<ArrayData> &data);

        BinaryArray(int64_t length, const std::shared_ptr<Buffer> &value_offsets,
                    const std::shared_ptr<Buffer> &data,
                    const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                    int64_t null_count = kUnknownNullCount, int64_t offset = 0);

    protected:
        // For subclasses such as StringArray
        BinaryArray() : BaseBinaryArray() {}
    };

    /// Concrete Array class for variable-size string (utf-8) data
    class TURBO_EXPORT StringArray : public BinaryArray {
    public:
        using TypeClass = StringType;

        explicit StringArray(const std::shared_ptr<ArrayData> &data);

        StringArray(int64_t length, const std::shared_ptr<Buffer> &value_offsets,
                    const std::shared_ptr<Buffer> &data,
                    const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                    int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        /// \brief Validate that this array contains only valid UTF8 entries
        ///
        /// This check is also implied by validate_full()
        turbo::Status validate_utf8() const;
    };

    /// Concrete Array class for large variable-size binary data
    class TURBO_EXPORT LargeBinaryArray : public BaseBinaryArray<LargeBinaryType> {
    public:
        explicit LargeBinaryArray(const std::shared_ptr<ArrayData> &data);

        LargeBinaryArray(int64_t length, const std::shared_ptr<Buffer> &value_offsets,
                         const std::shared_ptr<Buffer> &data,
                         const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                         int64_t null_count = kUnknownNullCount, int64_t offset = 0);

    protected:
        // For subclasses such as LargeStringArray
        LargeBinaryArray() : BaseBinaryArray() {}
    };

    /// Concrete Array class for large variable-size string (utf-8) data
    class TURBO_EXPORT LargeStringArray : public LargeBinaryArray {
    public:
        using TypeClass = LargeStringType;

        explicit LargeStringArray(const std::shared_ptr<ArrayData> &data);

        LargeStringArray(int64_t length, const std::shared_ptr<Buffer> &value_offsets,
                         const std::shared_ptr<Buffer> &data,
                         const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                         int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        /// \brief Validate that this array contains only valid UTF8 entries
        ///
        /// This check is also implied by validate_full()
        turbo::Status validate_utf8() const;
    };

    // ----------------------------------------------------------------------
    // BinaryView and StringView

    /// Concrete Array class for variable-size binary view data using the
    /// BinaryViewType::c_type struct to reference in-line or out-of-line string values
    class TURBO_EXPORT BinaryViewArray : public FlatArray {
    public:
        using TypeClass = BinaryViewType;
        using IteratorType = stl::ArrayIterator<BinaryViewArray>;
        using c_type = BinaryViewType::c_type;

        explicit BinaryViewArray(std::shared_ptr<ArrayData> data);

        BinaryViewArray(std::shared_ptr<DataType> type, int64_t length,
                        std::shared_ptr<Buffer> views, BufferVector data_buffers,
                        std::shared_ptr<Buffer> null_bitmap = nullptr,
                        int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        // For API compatibility with BinaryArray etc.
        [[nodiscard]] std::string_view get_view(int64_t i) const;

        [[nodiscard]] std::string get_string(int64_t i) const { return std::string{get_view(i)}; }

        [[nodiscard]] const auto &values() const { return data_->buffers[1]; }

        [[nodiscard]] const c_type *raw_values() const { return raw_values_; }

        std::optional<std::string_view> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        [[nodiscard]] IteratorType begin() const { return IteratorType(*this); }

        [[nodiscard]] IteratorType end() const { return IteratorType(*this, length()); }

    protected:
        using FlatArray::FlatArray;

        void set_data(std::shared_ptr<ArrayData> data) {
            FlatArray::set_data(std::move(data));
            raw_values_ = data_->GetValuesSafe<c_type>(1);
        }

        const c_type *raw_values_;
    };

    /// Concrete Array class for variable-size string view (utf-8) data using
    /// BinaryViewType::c_type to reference in-line or out-of-line string values
    class TURBO_EXPORT StringViewArray : public BinaryViewArray {
    public:
        using TypeClass = StringViewType;

        explicit StringViewArray(std::shared_ptr<ArrayData> data);

        using BinaryViewArray::BinaryViewArray;

        /// \brief Validate that this array contains only valid UTF8 entries
        ///
        /// This check is also implied by validate_full()
        turbo::Status validate_utf8() const;
    };

    // ----------------------------------------------------------------------
    // Fixed width binary

    /// Concrete Array class for fixed-size binary data
    class TURBO_EXPORT FixedSizeBinaryArray : public PrimitiveArray {
    public:
        using TypeClass = FixedSizeBinaryType;
        using IteratorType = stl::ArrayIterator<FixedSizeBinaryArray>;

        explicit FixedSizeBinaryArray(const std::shared_ptr<ArrayData> &data);

        FixedSizeBinaryArray(const std::shared_ptr<DataType> &type, int64_t length,
                             const std::shared_ptr<Buffer> &data,
                             const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                             int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        [[nodiscard]]  const uint8_t *get_value(int64_t i) const;

        [[nodiscard]] const uint8_t *value(int64_t i) const { return get_value(i); }

        [[nodiscard]]  std::string_view get_view(int64_t i) const {
            return std::string_view(reinterpret_cast<const char *>(get_value(i)), byte_width());
        }

        std::optional<std::string_view> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        [[nodiscard]] std::string get_string(int64_t i) const { return std::string(get_view(i)); }

        [[nodiscard]] int32_t byte_width() const { return byte_width_; }

        [[nodiscard]] const uint8_t *raw_values() const { return raw_values_ + data_->offset * byte_width_; }

        [[nodiscard]] IteratorType begin() const { return IteratorType(*this); }

        [[nodiscard]] IteratorType end() const { return IteratorType(*this, length()); }

    protected:
        void set_data(const std::shared_ptr<ArrayData> &data) {
            this->PrimitiveArray::set_data(data);
            byte_width_ =
                    turbo::checked_cast<const FixedSizeBinaryType &>(*type()).byte_width();
        }

        int32_t byte_width_;
    };

    /// @}

}  // namespace nebula
